#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Edit-Exif - 编辑图片的GPS信息

这个脚本用于直接编辑指定文件夹中所有图片的Exif GPS信息，
保留其他所有元数据不变，直接修改原始文件，保持文件创建时间等属性不变。
"""

import os
import sys
import argparse
from PIL import Image
from PIL.ExifTags import TAGS, GPSTAGS
from fractions import Fraction
import piexif
import shutil
from datetime import datetime


def get_decimal_from_dms(dms, ref):
    """
    将度分秒格式的坐标转换为十进制格式
    """
    degrees = dms[0]
    minutes = dms[1] / 60.0
    seconds = dms[2] / 3600.0
    
    decimal = degrees + minutes + seconds
    
    if ref in ['S', 'W']:
        decimal = -decimal
        
    return decimal


def get_dms_from_decimal(decimal):
    """
    将十进制格式的坐标转换为度分秒格式
    """
    is_negative = decimal < 0
    decimal = abs(decimal)
    
    degrees = int(decimal)
    minutes = int((decimal - degrees) * 60)
    seconds = (decimal - degrees - minutes/60) * 3600
    
    # 将秒转换为分数形式，以便存储在Exif中
    seconds_fraction = Fraction(seconds).limit_denominator(1000000)
    
    return ((degrees, 1), (minutes, 1), (seconds_fraction.numerator, seconds_fraction.denominator)), 'S' if is_negative and decimal != 0 else 'N'


def create_gps_exif(latitude, longitude, altitude=0):
    """
    创建GPS Exif数据
    """
    lat_dms, lat_ref = get_dms_from_decimal(latitude)
    lon_dms, lon_ref = get_dms_from_decimal(longitude)
    
    # 修正经度的参考方向
    if longitude < 0 and longitude != 0:
        lon_ref = 'W'
    else:
        lon_ref = 'E'
    
    # 创建GPS信息字典
    gps_ifd = {
        piexif.GPSIFD.GPSVersionID: (2, 2, 0, 0),
        piexif.GPSIFD.GPSLatitudeRef: lat_ref,
        piexif.GPSIFD.GPSLatitude: lat_dms,
        piexif.GPSIFD.GPSLongitudeRef: lon_ref,
        piexif.GPSIFD.GPSLongitude: lon_dms,
    }
    
    # 如果有海拔高度，也添加进去
    if altitude != 0:
        altitude_fraction = Fraction(altitude).limit_denominator(1000000)
        altitude_ref = 1 if altitude < 0 else 0  # 0表示海平面以上，1表示海平面以下
        gps_ifd[piexif.GPSIFD.GPSAltitudeRef] = altitude_ref
        gps_ifd[piexif.GPSIFD.GPSAltitude] = (altitude_fraction.numerator, altitude_fraction.denominator)
    
    return gps_ifd


def process_image(input_path, latitude, longitude, altitude=0, datetime=None):
    """
    处理单个图片，直接修改其GPS信息
    """
    try:
        # 对于不支持Exif的图片格式，跳过处理
        img_ext = os.path.splitext(input_path)[1].lower()
        if img_ext not in ['.jpg', '.jpeg', '.tiff', '.tif']:
            print(f"文件 {os.path.basename(input_path)} 不支持Exif，已跳过")
            return False
        
        # 读取原始图片的Exif数据
        try:
            exif_dict = piexif.load(input_path)
        except Exception as e:
            # 如果无法读取Exif，创建一个新的Exif字典
            exif_dict = {'0th': {}, '1st': {}, 'Exif': {}, 'GPS': {}, 'Interop': {}, 'thumbnail': None}
        
        # 创建新的GPS Exif数据
        gps_ifd = create_gps_exif(latitude, longitude, altitude)
        if datetime:
            exif_dict['Exif'][piexif.ExifIFD.DateTimeOriginal] = datetime
        
        # 更新Exif字典中的GPS部分
        exif_dict['GPS'] = gps_ifd
        
        # 将Exif数据转换为字节
        exif_bytes = piexif.dump(exif_dict)
        
        # 打开原始图片
        img = Image.open(input_path)
        
        # 保存图片，同时写入新的Exif数据（直接覆盖原文件）
        img.save(input_path, exif=exif_bytes)
        
        print(f"已处理: {os.path.basename(input_path)}")
        return True
    except Exception as e:
        print(f"处理 {os.path.basename(input_path)} 时出错: {str(e)}")
        return False


def process_directory(input_dir, latitude, longitude, altitude=0, recursive=False, datetime=None):
    """
    处理整个目录中的图片
    """
    # 支持的图片格式
    supported_formats = ['.jpg', '.jpeg', '.tiff', '.tif']
    
    # 计数器
    processed = 0
    failed = 0
    skipped = 0
    
    # 遍历目录
    for root, dirs, files in os.walk(input_dir):
        # 如果不是递归模式且不是顶级目录，则跳过
        if not recursive and root != input_dir:
            continue
        
        # 处理当前目录中的所有文件
        for file in files:
            file_ext = os.path.splitext(file)[1].lower()
            
            # 如果不是支持的图片格式，则跳过
            if file_ext not in supported_formats:
                skipped += 1
                continue
                
            input_path = os.path.join(root, file)
            
            # 处理图片
            if process_image(input_path, latitude, longitude, altitude, datetime):
                processed += 1
            else:
                failed += 1
    
    return processed, failed, skipped


def main():
    # 创建命令行参数解析器
    parser = argparse.ArgumentParser(description='直接修改图片的GPS Exif信息')
    parser.add_argument('input_dir', help='图片所在的文件夹路径')
    parser.add_argument('--lat', type=float, required=True, help='纬度值（十进制格式，如: 39.9042）')
    parser.add_argument('--lon', type=float, required=True, help='经度值（十进制格式，如: 116.4074）')
    parser.add_argument('--alt', type=float, default=0, help='海拔高度（米，默认为0）')
    parser.add_argument('--recursive', '-r', action='store_true', help='是否递归处理子文件夹')
    parser.add_argument('--datetime', type=str, help='照片的DateTimeOriginal信息（格式: YYYY:MM:DD HH:MM:SS）')
    
    # 解析命令行参数
    args = parser.parse_args()
    
    # 检查输入目录是否存在
    if args.datetime:
        try:
            datetime.strptime(args.datetime, '%Y:%m:%d %H:%M:%S')
        except ValueError:
            print(f"错误: 无效的DateTimeOriginal格式 '{args.datetime}'，应为 'YYYY:MM:DD HH:MM:SS'")
            return 1
    if not os.path.isdir(args.input_dir):
        print(f"错误: 输入目录 '{args.input_dir}' 不存在")
        return 1
    
    print(f"开始处理图片...")
    print(f"输入目录: {args.input_dir}")
    print(f"目标GPS坐标: 纬度 {args.lat}, 经度 {args.lon}, 海拔 {args.alt}米")
    print(f"递归处理子文件夹: {'是' if args.recursive else '否'}")
    print(f"警告: 此操作将直接修改原始文件，请确保已备份重要图片")
    print("-" * 50)
    
    # 处理目录
    processed, failed, skipped = process_directory(
        args.input_dir, args.lat, args.lon, args.alt, args.recursive, args.datetime
    )
    
    print("-" * 50)
    print(f"处理完成!")
    print(f"成功处理: {processed} 个文件")
    print(f"处理失败: {failed} 个文件")
    print(f"已跳过: {skipped} 个非图片文件")
    
    return 0


if __name__ == "__main__":
    sys.exit(main())